#!/usr/bin/env python

import rospy
from nav_msgs.msg import Odometry, Path
from geometry_msgs.msg import PoseStamped
import numpy as np
import acado
import time
import matplotlib.pyplot as plt

from trajectory_msgs.msg import  MultiDOFJointTrajectory,MultiDOFJointTrajectoryPoint
from sensor_msgs.msg import JointState
from nav_msgs.msg import Path
from geometry_msgs.msg import Transform, Twist
from ros_numpy import msgify,numpify
N=10 #timestep
NU=4
NOD=3
NX=16
NY=16
NYN=4
NW=16
NWN=12

target_position=np.array([3.0,0.0,0.0])



def odometry_callback(msg):
    # Process the odometry message and generate path points
    pose = msg.pose.pose  # Get the pose from the odometry message
    position = pose.position  # Get the position from the pose
    orientation = pose.orientation  # Get the orientation from the pose

    # Create a new pose stamped message
    pose_stamp = PoseStamped()
    pose_stamp.pose.position = position
    pose_stamp.pose.orientation = orientation
    
    
    x0=np.zeros(NX).reshape(1,NX)
    u=np.zeros([N,NU])

    Q=np.eye(NW)
    Q=np.diag([1e-8,1e-8,1e-8,1,1,1,1,100,10,10,100,100,1,5,5,5e5])
    W=np.transpose(np.tile(Q,N))
    WN=Q[:-4,:-4]
    y=np.zeros([N,NY])

    OD=np.ones([N+1,NOD])
    for i in range(N):
        OD[i,0:3]=np.array([3.3,0,0.67])+np.array([0.0,0,0])*i
    
    x0[0,0]=position.x #px
    x0[0,1]=position.y#py
    x0[0,2]=position.z #pz
    x0[0,3]=1.0#qw
    x0[0,4]=0.0#qx
    x0[0,5]=0.0#qy
    x0[0,6]=0.0#qz
    x0[0,7]=0.0#vx
    x0[0,8]=0.0#vy
    x0[0,9]=0.0#vz
    x0[0,10]=1.0#pxx
    x0[0,11]=0.0#pxy
    x0[0,12]=0.0#pxz
    x0[0,13]=1.0#pyy
    x0[0,14]=0.0#pyz
    x0[0,15]=1.0#pzz
    
    X=np.tile(x0.reshape(NX,1),N+1).T
    
    X,u=acado.mpc(0,1,x0,X,u,OD,y,W,WN,1)
    x0=X[1,:].reshape(1,NX)
    
    u=np.zeros([N,NU])

    xlog=[]
    tracelog=[]

    t0=time.time()
    
        # Publish the pose stamped message
    path_msg = Path()
    path_msg.header = msg.header  # Set the header of the path message
    
    traj_msg=MultiDOFJointTrajectory()

    for i in range(20):
        
    
        #print("iter: "+str(i))
        X,u=acado.mpc(0,1,x0,X,u,OD,y,W,WN,1)
        x0=X[1,:].reshape(1,NX)
        P=np.array([[x0[0,10],x0[0,11],x0[0,12]],
                    [x0[0,11],x0[0,13],x0[0,14]],
                    [x0[0,12],x0[0,14],x0[0,15]]])
        
        
        #tracelog.append(x0[0,10]+x0[0,13]+x0[0,15])
        tracelog.append(np.linalg.det(P))
        
        xlog.append(x0)
        xlog_arr=np.vstack(xlog)
        

        pose_stamp=PoseStamped()
        pose_stamp.pose.position.x=x0[0,0]
        pose_stamp.pose.position.y=x0[0,1]
        pose_stamp.pose.position.z=x0[0,2]

        path_msg.poses.append(pose_stamp)
        
        #meanwhile pub the trajectory msg
        traj_point=MultiDOFJointTrajectoryPoint()
        transform=Transform()
        transform.translation.x=x0[0,0]
        transform.translation.y=x0[0,1]
        transform.translation.z=x0[0,2]
        
        twist=Twist()
        twist.linear.x=x0[0,7]
        twist.linear.y=x0[0,8]
        twist.linear.z=x0[0,9]
        
        traj_point.transforms.append(transform)
        traj_point.velocities.append(twist)
        traj_msg.points.append(traj_point)                        
        
        

    path_pub.publish(path_msg)  # Publish the path message
    traj_pub.publish(traj_msg)
    
def mav_pose_callback(msg):
    print(target_position)
    
def target_pose_callback(msg):
    target_pose=numpify(msg.pose)
    
    target_position=target_pose[:3,3]
    print(target_position)
    
    
    

if __name__ == '__main__':
    rospy.init_node('pos_planning_node')  # Initialize the ROS node
    rospy.Subscriber('/bebop2/ground_truth/odometry', Odometry, odometry_callback)  # Subscribe to the "Odometry" topic
    rospy.Subscriber('/slot/mav/pose',PoseStamped,mav_pose_callback)
    rospy.Subscriber('/slot/target/pose',PoseStamped,target_pose_callback)
    path_pub = rospy.Publisher('/path', Path, queue_size=10)  # Publish the "Path" topic
    traj_pub=rospy.Publisher('/position_traj',MultiDOFJointTrajectory,queue_size=10)
    rospy.spin()  # Keep the node running until it is stopped